#include <algorithm>
#include <cmath>
#include <iostream>
#include <string>
#include <vector>

#include "example_utils.hpp"
#include "oneapi/dnnl/dnnl.hpp"

using namespace dnnl;
using namespace std;

using tag = memory::format_tag;
using dt = memory::data_type;
void spatialtf_example(dnnl::engine::kind engine_kind) {
    // Create execution dnnl::engine.
    dnnl::engine engine(engine_kind, 0);

    // Create dnnl::stream.
    dnnl::stream engine_stream(engine);

    // Tensor dimensions.
    const memory::dim B = 1, // batch size
            IC = 1, // channels
            IH = 5, // height
            IW = 5; // width

    memory::dims src_dims =  {B,IC,IH,IW};
    memory::dims dst_dims =  {B,IH,IW,2};

    // Allocate buffer.
    std::vector<float> src_data(product(src_dims));
    for (size_t i = 0; i < src_data.size(); i++) {
        src_data[i] = 1.0 / i;
    }

    const memory::dim thetaHeight = 2,thetaWidth = 3;

    memory::dims theta_dims =  {1,thetaHeight,thetaWidth};
    std::vector<float> theta_data(thetaHeight * thetaWidth);

    for(int i = 0; i < thetaHeight * thetaWidth; i++)
    {
        theta_data[i] = 1.0;
    }

    // Create src memory descriptor and memory object.
    auto src0_md = memory::desc(src_dims, dt::f32, tag::nchw);
    auto src1_md = memory::desc(theta_dims, dt::f32, tag::abc);
    auto dst_md = memory::desc(dst_dims, dt::f32, tag::nchw);
    auto src0_mem = memory(src0_md, engine);
    auto theta_mem = memory(src1_md, engine);
    auto dst_mem = memory(dst_md, engine);
    std::vector<float> dst_data(product(dst_dims));

//    write_to_dnnl_memory(src_data.data(), src0_mem);
    // Write data to memory object's handle.
    write_to_dnnl_memory(theta_data.data(), theta_mem);

    auto spatialtf_pd = spatialtf_grid_generator_forward::primitive_desc(engine,
            prop_kind::forward_training, src0_md, src1_md, dst_md);


    // Create the primitive.
    auto spatialtf_prim = spatialtf_grid_generator_forward(spatialtf_pd);

    // Primitive arguments. Set up in-place execution by assigning src as DST.
    std::unordered_map<int, memory> spatialtf_args;

    spatialtf_args.insert({DNNL_ARG_SRC_0, src0_mem});
    spatialtf_args.insert({DNNL_ARG_SRC_1, theta_mem});
    spatialtf_args.insert({DNNL_ARG_DST, dst_mem});

    // Primitive execution.
    spatialtf_prim.execute(engine_stream, spatialtf_args);

    // Wait for the computation to finalize.
    engine_stream.wait();

    // Read data from memory object's handle.
    read_from_dnnl_memory(dst_data.data(), dst_mem);
    for (size_t i = 0; i < dst_data.size(); i++) {
        std::cout << dst_data[i] << " ";
    }

}
int main(int argc, char **argv) {
    return handle_example_errors(
            spatialtf_example, parse_engine_kind(argc, argv));
}